import json
import os
import sys
import time
from typing import Dict, List, Optional
import re

# 确保项目根目录在Python路径中
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入提示词量化评估中的SparkModelEvaluator类
try:
    # 使用绝对导入路径
    from 提示词量化评估.spark_evaluator import SparkModelEvaluator
    print("成功导入SparkModelEvaluator类")
except ImportError as e:
    print(f"警告: 无法导入SparkModelEvaluator类，评估功能将不可用。错误信息: {str(e)}")
    SparkModelEvaluator = None

# 如果直接导入失败，尝试添加具体的提示词量化评估目录到Python路径
try:
    if SparkModelEvaluator is None:
        evaluation_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "提示词量化评估")
        if evaluation_dir not in sys.path:
            sys.path.append(evaluation_dir)
        from spark_evaluator import SparkModelEvaluator
        print("通过添加评估目录到Python路径成功导入SparkModelEvaluator类")
except ImportError as e:
    if SparkModelEvaluator is None:
        print(f"再次尝试导入失败: {str(e)}")

class PromptOptimizer:
    def __init__(self):
        self.evaluator = SparkModelEvaluator() if SparkModelEvaluator else None # 构造SparkModelEvalutaor类的实例，如果该类存在的话，否则返回None
        self.optimization_history = []
        self.script_dir = os.path.dirname(os.path.abspath(__file__))            # 获取当前脚本所在目录，用于保存优化结果
        if self.evaluator:      # 如果evaluator存在，则可以直接调用星火模型进行优化
            print("已成功初始化SparkModelEvaluator，可以使用星火大模型进行提示词优化和评估")
        else:
            print("注意: SparkModelEvaluator初始化失败，将使用本地CRISP框架进行提示词优化")
    def extract_crisp_elements(self, long_text: str) -> Dict[str, str]:
        """
        从长文本中提取CRISP框架的各个元素
        参数:long_text: 用户输入的长文本提示词
        返回:包含CRISP各元素的字典
        """
        crisp_elements = {
            "context": "",   # 上下文
            "role": "",       # 角色
            "instruction": "", # 指令
            "style": "",      # 风格
            "purpose": ""      # 目的
        }
        # 1. 尝试识别角色（Role）
        role_patterns = [
            r"你是(.*?)\.",
            r"假设你是(.*?)\.",
            r"请以(.*?)的身份",
            r"请扮演(.*?)"
        ]
        for pattern in role_patterns:
            match = re.search(pattern, long_text)   # 利用re组件，从长文本中搜索角色模式
            if match:
                crisp_elements["role"] = match.group(1).strip()
                long_text = re.sub(pattern, "", long_text)
                break
        if not crisp_elements["role"]:  # 如果没有明确识别出角色，设置默认角色
            crisp_elements["role"] = "专业助手"
        # 2. 尝试识别目的（Purpose）
        purpose_patterns = [
            r"目的是(.*?)\.",
            r"旨在(.*?)\.",
            r"为了(.*?)\.",
            r"目标是(.*?)\."
        ]
        for pattern in purpose_patterns:
            match = re.search(pattern, long_text)
            if match:
                crisp_elements["purpose"] = match.group(1).strip()
                long_text = re.sub(pattern, "", long_text)
                break
        # 3. 尝试识别风格（Style）
        style_patterns = [
            r"风格(.*?)",
            r"以(.*?)的风格",
            r"请用(.*?)的方式",
            r"要求(.*?)的表达"
        ]
        for pattern in style_patterns:
            match = re.search(pattern, long_text)
            if match:
                crisp_elements["style"] = match.group(1).strip()
                long_text = re.sub(pattern, "", long_text)
                break
        # 如果没有明确识别出风格，设置默认风格
        if not crisp_elements["style"]:
            crisp_elements["style"] = "简洁明了"
        # 4. 提取指令（Instruction）- 通常是核心任务描述
        # 这里简化处理，假设剩余的大部分文本都是指令
        instruction_keywords = ["请", "需要", "要求", "希望", "烦请", "能否"]
        for keyword in instruction_keywords:
            idx = long_text.find(keyword)
            if idx != -1:
                instruction_part = long_text[idx:].strip()  # 从关键词开始到文本结束作为指令
                period_idx = instruction_part.find("。")    # 尝试分割上下文和指令
                if period_idx != -1 and period_idx < len(instruction_part) * 0.3:  # 如果句号在前30%的位置，可能是上下文
                    crisp_elements["context"] = instruction_part[:period_idx+1].strip()
                    crisp_elements["instruction"] = instruction_part[period_idx+1:].strip()
                else:
                    crisp_elements["instruction"] = instruction_part
                break
        if not crisp_elements["instruction"]:   # 如果没有找到明确的指令关键词，使用剩余文本作为指令
            crisp_elements["instruction"] = long_text.strip()
        return crisp_elements
    
    def generate_crisp_prompt(self, crisp_elements: Dict[str, str]) -> str:
        """
        根据CRISP元素生成结构化提示词
        参数:
            crisp_elements: 包含CRISP各元素的字典
        返回:
            结构化的CRISP提示词
        """
        # 构建CRISP结构化提示词
        crisp_prompt = []
        # 1. 添加角色（Role）
        if crisp_elements["role"]:
            crisp_prompt.append(f"角色: 你是{crisp_elements['role']}")
        # 2. 添加上下文（Context）
        if crisp_elements["context"]:
            crisp_prompt.append(f"上下文: {crisp_elements['context']}")
        # 3. 添加目的（Purpose）
        if crisp_elements["purpose"]:
            crisp_prompt.append(f"目的: {crisp_elements['purpose']}")
        # 4. 添加风格（Style）
        if crisp_elements["style"]:
            crisp_prompt.append(f"风格: 请以{crisp_elements['style']}的风格")
        # 5. 添加指令（Instruction）
        if crisp_elements["instruction"]:
            crisp_prompt.append(f"指令: {crisp_elements['instruction']}")
        # 组合所有元素为一个完整的提示词
        return "\n".join(crisp_prompt)
    
    def optimize_with_spark(self, long_text: str) -> str:
        """
        使用讯飞星火大模型直接优化提示词
        参数:
            long_text: 用户输入的长文本提示词
        返回:
            优化后的提示词
        """
        if not self.evaluator:
            print("错误: SparkModelEvaluator不可用，无法使用星火大模型进行优化")
            crisp_elements = self.extract_crisp_elements(long_text) # 如果evaluator不可用，使用默认的CRISP框架优化
            return self.generate_crisp_prompt(crisp_elements)
        print("正在使用讯飞星火大模型优化提示词...")
        # 构建用于优化提示词的提示
        optimize_prompt = f"""
        请你作为一名提示词优化专家，帮我优化下面的长文本提示词，使其更简洁、精准、有效。
        优化要求：
        1. 保留核心信息和意图
        2. 提高提示词的清晰度和明确性
        3. 使提示词更符合大模型的理解习惯
        4. 保持简洁，但不要遗漏重要信息
        5. 优化后的提示词应保持原始意图不变
        原始提示词：
        {long_text}
        请直接返回优化后的提示词，不要添加任何额外的解释或说明。
        """
        # 调用星火模型进行优化
        optimized_prompt, _ = self.evaluator.call_spark_model(optimize_prompt, temperature=0.7) # 调用星火模型进行优化
        # 如果优化失败，使用默认的CRISP框架优化
        if not optimized_prompt:
            print("星火大模型优化失败，使用默认的CRISP框架优化")
            crisp_elements = self.extract_crisp_elements(long_text)
            optimized_prompt = self.generate_crisp_prompt(crisp_elements)
        return optimized_prompt.strip()
    
    def optimize_prompt(self, long_text: str, use_spark: bool = True) -> Dict[str, str]:
        """
        优化长文本提示词，可以选择使用讯飞星火大模型或CRISP框架
        参数:
            long_text: 用户输入的长文本提示词
            use_spark: 是否使用讯飞星火大模型进行优化
        返回:
            包含原始提示词、提取的CRISP元素和优化后提示词的字典
        """
        print(f"正在优化提示词: {long_text[:50]}{'...' if len(long_text) > 50 else ''}")
        crisp_elements = self.extract_crisp_elements(long_text)     # 提取CRISP元素（无论使用哪种优化方法，都提取CRISP元素用于记录）
        # 根据use_spark参数和evaluator是否可用选择优化方法
        if use_spark and self.evaluator:
            optimized_prompt = self.optimize_with_spark(long_text)
            optimization_method = "spark_model"
        else:
            # 使用传统的CRISP框架优化
            optimized_prompt = self.generate_crisp_prompt(crisp_elements)
            optimization_method = "crisp_framework"
        # 记录优化历史
        optimization_result = {
            "original_prompt": long_text,
            "crisp_elements": crisp_elements,
            "optimized_prompt": optimized_prompt,
            "optimization_method": optimization_method,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        self.optimization_history.append(optimization_result)
        return optimization_result
    
    def compare_prompts(self, original_prompt: str, optimized_prompt: str) -> Optional[Dict]:
        """
        比较原始提示词和优化后提示词的效果
        参数:
            original_prompt: 原始长文本提示词
            optimized_prompt: 优化后的CRISP框架提示词
        返回:
            包含比较结果的字典，如果评估器不可用则返回None
        """
        if not self.evaluator:
            print("评估器不可用，无法比较提示词效果")
            return None
        
        print("正在比较原始提示词和优化后提示词的效果...")
        
        # 评估原始提示词
        print("评估原始提示词...")
        original_result = self.evaluator.evaluate_prompt(original_prompt)
        
        # 评估优化后的提示词
        print("评估优化后的提示词...")
        optimized_result = self.evaluator.evaluate_prompt(optimized_prompt)
        
        # 构建比较结果
        comparison = {
            "original_prompt": original_prompt,
            "optimized_prompt": optimized_prompt,
            "original_scores": original_result["avg_scores"] if original_result else None,
            "optimized_scores": optimized_result["avg_scores"] if optimized_result else None,
            "improvement": None
        }
        
        # 计算改进度
        if original_result and optimized_result:
            original_total = original_result["avg_scores"]["total"]
            optimized_total = optimized_result["avg_scores"]["total"]
            comparison["improvement"] = round(optimized_total - original_total, 2)
            
        return comparison
    
    def save_optimization_result(self, result: Dict, filename: str = None):
        """
        保存优化结果到文件
        
        参数:
            result: 优化结果字典
            filename: 保存的文件名，默认为None（自动生成）
        """
        # 如果没有提供文件名，根据提示词和时间戳生成唯一文件名
        if filename is None:
            # 获取当前时间戳，确保唯一性
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            
            # 从结果中提取提示词信息
            if result and "original_prompt" in result:
                prompt_preview = result["original_prompt"][:10] if result["original_prompt"] else "unknown"
            else:
                prompt_preview = "unknown"
            
            # 替换文件名中的非法字符
            prompt_preview = "".join(c for c in prompt_preview if c.isalnum() or c in ("_", "-", " ")).strip()
            
            # 构建文件名
            filename = f"optimization_result_{prompt_preview}_{timestamp}.json"
        
        # 如果提供的文件名不是绝对路径，则将其放在脚本所在目录下
        if not os.path.isabs(filename):
            filename = os.path.join(self.script_dir, filename)
        
        # 保存结果
        with open(filename, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        
        print(f"优化结果已保存到 {filename}")
    
    def print_optimization_result(self, result: Dict):
        """
        打印优化结果
        
        参数:
            result: 优化结果字典
        """
        print("\n" + "="*60)
        print("原始提示词:")
        print(result["original_prompt"][:200] + ("..." if len(result["original_prompt"]) > 200 else ""))
        
        print("\n提取的CRISP元素:")
        for key, value in result["crisp_elements"].items():
            print(f"  {key}: {value}")
        
        print("\n优化后的提示词:")
        print(result["optimized_prompt"])
        
        print("="*60 + "\n")
    
    def print_comparison_result(self, comparison: Dict):
        """
        打印提示词比较结果
        
        参数:
            comparison: 比较结果字典
        """
        if not comparison:
            return
        
        print("\n" + "="*60)
        print("提示词效果比较结果")
        print("="*60)
        
        print(f"原始提示词总分: {comparison['original_scores']['total'] if comparison['original_scores'] else 'N/A'}")
        print(f"优化后提示词总分: {comparison['optimized_scores']['total'] if comparison['optimized_scores'] else 'N/A'}")
        
        if comparison["improvement"] is not None:
            improvement_str = f"+{comparison['improvement']}" if comparison["improvement"] > 0 else str(comparison["improvement"])
            print(f"改进度: {improvement_str}")
            
            if comparison["improvement"] > 0:
                print("✓ 优化后的提示词效果更好！")
            elif comparison["improvement"] < 0:
                print("⚠️ 优化后的提示词效果有所下降。")
            else:
                print("优化后的提示词效果与原始提示词相当。")
        
        print("="*60 + "\n")

def main():
    """主函数"""
    optimizer = PromptOptimizer()
    
    print("欢迎使用长文本提示词优化工具")
    print("此工具可以将长文本提示词转换为更精准、更有效的提示词，支持使用CRISP框架或讯飞星火大模型进行优化。")
    
    while True:
        print("\n请选择操作:")
        print("1. 使用CRISP框架优化单个提示词")
        print("2. 使用讯飞星火大模型优化单个提示词")
        print("3. 比较优化前后的提示词效果")
        print("4. 退出")
        
        choice = input("请输入选择 (1/2/3/4): ")
        
        if choice == "1":
            long_text = input("请输入要优化的长文本提示词: ")
            result = optimizer.optimize_prompt(long_text, use_spark=False)
            optimizer.print_optimization_result(result)
            
            # 询问是否保存结果
            save_choice = input("是否保存优化结果？(y/n): ").lower()
            if save_choice == "y":
                optimizer.save_optimization_result(result)
        
        elif choice == "2":
            if not optimizer.evaluator:
                print("错误: SparkModelEvaluator不可用，无法使用星火大模型进行优化。将使用CRISP框架代替。")
                use_spark = False
            else:
                use_spark = True
            
            long_text = input("请输入要优化的长文本提示词: ")
            result = optimizer.optimize_prompt(long_text, use_spark=use_spark)
            optimizer.print_optimization_result(result)
            
            # 询问是否保存结果
            save_choice = input("是否保存优化结果？(y/n): ").lower()
            if save_choice == "y":
                optimizer.save_optimization_result(result)
        
        elif choice == "3":
            if not optimizer.evaluator:
                print("错误: 无法导入评估模块，无法比较提示词效果。")
                continue
            
            long_text = input("请输入要优化并比较的长文本提示词: ")
            
            # 分别使用两种方法进行优化
            print("\n使用CRISP框架优化...")
            crisp_result = optimizer.optimize_prompt(long_text, use_spark=False)
            print("\n使用讯飞星火大模型优化...")
            spark_result = optimizer.optimize_prompt(long_text, use_spark=True)
            
            # 打印两种优化结果
            print("\n========= CRISP框架优化结果 ========")
            optimizer.print_optimization_result(crisp_result)
            print("\n========= 讯飞星火大模型优化结果 ========")
            optimizer.print_optimization_result(spark_result)
            
            # 比较两种优化结果
            print("\n比较CRISP框架优化与原始提示词:")
            crisp_comparison = optimizer.compare_prompts(long_text, crisp_result["optimized_prompt"])
            optimizer.print_comparison_result(crisp_comparison)
            
            print("\n比较讯飞星火大模型优化与原始提示词:")
            spark_comparison = optimizer.compare_prompts(long_text, spark_result["optimized_prompt"])
            optimizer.print_comparison_result(spark_comparison)
            # 保存比较结果
            combined_result = {
                "original_prompt": long_text,
                "crisp_optimization_result": crisp_result,
                "spark_optimization_result": spark_result,
                "crisp_comparison_result": crisp_comparison,
                "spark_comparison_result": spark_comparison
            }
            optimizer.save_optimization_result(combined_result)
        elif choice == "4":
            print("感谢使用，再见！")
            break
        else:
            print("无效的选择，请重新输入。")

if __name__ == "__main__":
    main()